///
/// @file SimpleRegisterAllocator.cpp
/// @brief 简单或朴素的寄存器分配器
/// @author zenglj (zenglj@live.com)
/// @version 1.0
/// @date 2024-09-29
///
/// @copyright Copyright (c) 2024
///
/// @par 修改日志:
/// <table>
/// <tr><th>Date       <th>Version <th>Author  <th>Description
/// <tr><td>2024-09-29 <td>1.0     <td>zenglj  <td>新建
/// </table>
///
#include <algorithm>
#include "SimpleRegisterAllocator.h"
#include "Common.h"
#include "ILocArm64.h"

// 构造函数
SimpleRegisterAllocator::SimpleRegisterAllocator()
{}

// 整型寄存器分配
int SimpleRegisterAllocator::AllocateInt(Value * var, int32_t no)
{
    if (var && (var->getLoadRegId() != -1))
        return var->getLoadRegId();
    int32_t regno = -1;

    // 如果指定了特定寄存器且可用，则使用该寄存器
    if ((no != -1) && !regBitmap.test(no)) {
        regno = no;
    } else {
        for (int i = 0; i < PlatformArm64::maxUsableRegNum; ++i) {
            int idx = (currentReg + i) % PlatformArm64::maxUsableRegNum;
            if (!regBitmap.test(idx)) {
                regno = idx;
                currentReg = (idx + 1) % PlatformArm64::maxUsableRegNum;
                break;
            }
        }
    }
    if (regno == -1) {
        Value * spillVar = regValues.front();
        regno = spillVar->getLoadRegId();
        spillVar->setLoadRegId(-1);
        regValues.erase(regValues.begin());
        currentReg = (regno + 1) % PlatformArm64::maxUsableRegNum;
    }
    if (var) {
        var->setLoadRegId(regno);
        regValues.push_back(var);
    }
    bitmapSet(regno);
    return regno;
}

void SimpleRegisterAllocator::AllocateInt(int32_t no)
{
    if (regBitmap.test(no))
        freeInt(no);
    bitmapSet(no);
}

void SimpleRegisterAllocator::freeInt(Value * var)
{
    if (var && (var->getLoadRegId() != -1)) {
        regBitmap.reset(var->getLoadRegId());
        regValues.erase(std::find(regValues.begin(), regValues.end(), var));
        var->setLoadRegId(-1);
    }
}

void SimpleRegisterAllocator::freeInt(int32_t no)
{
    if (no == -1)
        return;
    regBitmap.reset(no);
    auto pIter = std::find_if(regValues.begin(), regValues.end(), [=](auto val) { return val->getLoadRegId() == no; });
    if (pIter != regValues.end()) {
        (*pIter)->setLoadRegId(-1);
        regValues.erase(pIter);
    }
}

void SimpleRegisterAllocator::bitmapSet(int32_t no)
{
    regBitmap.set(no);
    usedBitmap.set(no);
}

// 浮点寄存器分配
int SimpleRegisterAllocator::AllocateFloat(Value * var, int32_t no)
{
    if (var && (var->getLoadRegId() != -1))
        return var->getLoadRegId();
    int32_t regno = -1;
    if ((no != -1) && !floatRegBitmap.test(no)) {
        regno = no;
    } else {
        for (int i = 0; i < PlatformArm64::maxUsableFloatNum; ++i) {
            int idx = (currentFloatReg + i) % PlatformArm64::maxUsableFloatNum;
            if (!floatRegBitmap.test(idx)) {
                regno = idx;
                currentFloatReg = (idx + 1) % PlatformArm64::maxUsableFloatNum;
                break;
            }
        }
    }
    if (regno == -1) {
        // 这里说明没有可用寄存器，应该做溢出处理或报错
        minic_log(LOG_ERROR,
                  "AllocateFloat failed: no available float register for variable %s",
                  var ? var->getName().c_str() : "(null)");
        abort(); // 立即终止，防止生成非法代码
        Value * spillVar = floatRegValues.front();
        regno = spillVar->getLoadRegId();
        spillVar->setLoadRegId(-1);
        floatRegValues.erase(floatRegValues.begin());
        currentFloatReg = (regno + 1) % PlatformArm64::maxUsableFloatNum;
    }
    if (var) {
        var->setLoadRegId(regno);
        floatRegValues.push_back(var);
    }
    floatBitmapSet(regno);
    return regno;
}

void SimpleRegisterAllocator::AllocateFloat(int32_t no)
{
    if (floatRegBitmap.test(no))
        freeFloat(no);
    floatBitmapSet(no);
}

void SimpleRegisterAllocator::freeFloat(Value * var)
{
    if (var && (var->getLoadRegId() != -1)) {
        floatRegBitmap.reset(var->getLoadRegId());
        floatRegValues.erase(std::find(floatRegValues.begin(), floatRegValues.end(), var));
        var->setLoadRegId(-1);
    }
}

void SimpleRegisterAllocator::freeFloat(int32_t no)
{
    if (no == -1)
        return;
    floatRegBitmap.reset(no);
    auto pIter =
        std::find_if(floatRegValues.begin(), floatRegValues.end(), [=](auto val) { return val->getLoadRegId() == no; });
    if (pIter != floatRegValues.end()) {
        (*pIter)->setLoadRegId(-1);
        floatRegValues.erase(pIter);
    }
}

void SimpleRegisterAllocator::floatBitmapSet(int32_t no)
{
    if (no < 0 || no >= PlatformArm64::maxUsableFloatNum)
        return; // 防御
    floatRegBitmap.set(no);
    floatUsedBitmap.set(no);
}

// 统一接口
int SimpleRegisterAllocator::Allocate(bool is_float, Value * var, int32_t no)
{
    return is_float ? AllocateFloat(var, no) : AllocateInt(var, no);
}

void SimpleRegisterAllocator::AllocateReg(int32_t no, bool is_float)
{
    if (is_float)
        AllocateFloat(no);
    else
        AllocateInt(no);
}

void SimpleRegisterAllocator::free(Value * var, bool is_float)
{
    if (is_float)
        freeFloat(var);
    else
        freeInt(var);
}

void SimpleRegisterAllocator::free(int32_t no, bool is_float)
{
    if (is_float)
        freeFloat(no);
    else
        freeInt(no);
}

// 保存所有活跃的寄存器到栈上，返回压栈的字节数
int SimpleRegisterAllocator::saveActiveRegisters(ILocArm64 & iloc)
{
    std::vector<int> activeRegs;

    // 收集所有活跃的整型寄存器（x8-x28）
    for (int i = 8; i < PlatformArm64::maxUsableRegNum; i++) {
        if (regBitmap.test(i)) {
            activeRegs.push_back(i);
        }
    }

    if (activeRegs.empty()) {
        return 0;
    }

    // 计算需要的栈空间（16字节对齐）
    int space = activeRegs.size() * 8;
    space = (space + 15) & ~15; // 16字节对齐

    // 分配栈空间
    iloc.inst("sub", "sp", "sp", "#" + std::to_string(space));

    // 保存寄存器到栈上
    for (size_t i = 0; i < activeRegs.size(); i++) {
        iloc.inst("str", PlatformArm64::regName[activeRegs[i]], "[sp, #" + std::to_string(i * 8) + "]");
    }

    return space;
}

// 从栈上恢复所有保存的寄存器
void SimpleRegisterAllocator::restoreActiveRegisters(ILocArm64 & iloc, int stackSize)
{
    if (stackSize == 0) {
        return;
    }

    std::vector<int> activeRegs;

    // 收集所有活跃的整型寄存器（x8-x28）
    for (int i = 8; i < PlatformArm64::maxUsableRegNum; i++) {
        if (regBitmap.test(i)) {
            activeRegs.push_back(i);
        }
    }

    // 恢复寄存器
    for (size_t i = 0; i < activeRegs.size(); i++) {
        iloc.inst("ldr", PlatformArm64::regName[activeRegs[i]], "[sp, #" + std::to_string(i * 8) + "]");
    }

    // 释放栈空间
    iloc.inst("add", "sp", "sp", "#" + std::to_string(stackSize));
}